# coding: utf-8


import cv2
import os
import numpy as np
from metrics import *
import pandas as pd
from PIL import Image
Image.MAX_IMAGE_PIXELS = None
from tqdm import tqdm

#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++#
#                   计算 x265 编码器的性能指标， 图像数据集： CLIC、Kodak、Tecnick
#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++#

for i in range(1, 6, 1):
    print(i)
    # origin_dataset_folder = r"D:\video-communication-dataset\图像编解码测试\CLIC\professional_valid"
    # origin_dataset_folder = r"D:\video-communication-dataset\图像编解码测试\Tecnick\TESTIMAGES\RGB\RGB_OR_1200x1200"
    origin_dataset_folder = r"D:\video-communication-dataset\图像编解码测试\kodak\archive"

    # x265_image_folder = r"G:\phd_data\语义通信实验\HM16.26\image\CLIC\ldpc+16qam+snr%d" % (i)
    # x265_image_folder = r"G:\phd_data\语义通信实验\HM16.26\image\Tecnick\ldpc+16qam+snr%d" % (i)
    x265_image_folder = r"G:\phd_data\语义通信实验\HM16.26\image\Kodak\ldpc+16qam+snr%d" % (i)

    # x265_mkv_folder = r"G:\phd_data\语义通信实验\HM16.26\image\CLIC\CLIC_encode_bits"
    # x265_mkv_folder = r"G:\phd_data\语义通信实验\HM16.26\image\Tecnick\Tecnick_encode_bits"
    x265_mkv_folder = r"G:\phd_data\语义通信实验\HM16.26\image\Kodak\Kodak_encode_bits"

    # xlsx_save_path = r"G:\phd_data\语义通信实验\HM16.26\image\CLIC\ldpc_qam_awgn_snr%d.xlsx" % (i)
    # xlsx_save_path = r"G:\phd_data\语义通信实验\HM16.26\image\Tecnick\ldpc_qam_awgn_snr%d.xlsx" % (i)
    xlsx_save_path = r"G:\phd_data\语义通信实验\HM16.26\image\Kodak\ldpc_qam_awgn_snr%d.xlsx" % (i)

    qps = [str(i) for i in range(15, 52, 2)]

    data = {
        f"{i}": dict() for i in qps
    }

    image_names = os.listdir(origin_dataset_folder)

    for qp in qps:
        print(f"\nprocess qp {qp}")
        img_names, bpp_list, ssim_list, msssim_list, psnr_list = [], [], [], [], []
        for image_name in tqdm(image_names):
            image_path = os.path.join(origin_dataset_folder, image_name)
            hm_image_path = os.path.join(x265_image_folder, qp, image_name.replace(".png", "_dec.png"))
            mkv_image_path = os.path.join(x265_mkv_folder, qp, image_name.replace(".png", ".bin"))

            origin_image = Image.open(image_path).convert('RGB')
            H, W = origin_image.size
            x265_image = Image.open(hm_image_path).convert('RGB')
            # BPP
            bpp_val = compute_bpp(mkv_image_path, height=H, width=W)
            # SSIM
            ssim_val = compute_ssim(origin_image, x265_image)
            # MS-SSIM
            msssim_val = compute_msssim(origin_image, x265_image)
            # PSNR
            psnr_val = compute_psnr(origin_image, x265_image)
            # append
            img_names.append(image_name)
            bpp_list.append(bpp_val)
            ssim_list.append(ssim_val)
            msssim_list.append(msssim_val)
            psnr_list.append(psnr_val)
        data[qp]["Image name"] = img_names
        data[qp]["BPP"] = bpp_list
        data[qp]["SSIM"] = ssim_list
        data[qp]["MS-SSIM"] = msssim_list
        data[qp]["PSNR"] = psnr_list
        # append
        data[qp]["Image name"].append("average:")
        data[qp]["BPP"].append(np.array(bpp_list).mean())
        data[qp]["SSIM"].append(np.array(ssim_list).mean())
        data[qp]["MS-SSIM"].append(np.array(msssim_list).mean())
        data[qp]["PSNR"].append(np.array(psnr_list).mean())
    # save to excel:
    # qp-i: image name, bpp, ssim ms-ssim, psnr
    with pd.ExcelWriter(xlsx_save_path, engine="openpyxl") as writer:
        for quality, metrics in data.items():
            df = pd.DataFrame(metrics)
            df.to_excel(writer, sheet_name=quality, index=False)

